from helper import utils
import paddle
from paddle import nn
from paddle import io
from visualdl import LogWriter
import numpy as np
import pandas as pd
import json
import os
import random
from time import time
from sklearn.metrics import (roc_auc_score, average_precision_score, f1_score, roc_curve, confusion_matrix, 
                             precision_score, recall_score, auc, mean_squared_error)
from argparse import ArgumentParser

from double_towers import MolTransModel
from preprocess import drug_encoder, target_encoder


#CUDA_VISIBLE_DEVICES=0

#从外部读取蛋白质序列文件
protein_data = pd.read_csv('protein.csv')
protein_sequences = protein_data['Sequence']


#定义输入
d1='CC1=C2C=C(C=CC2=NN1)C3=CC(=CN=C3)OCC(CC4=CC=CC=C4)N'
#t1='MLKFKYGARNPLDAGAAEPIASRASRLNLFFQGKPPFMTQQQMSPLSREGILDALFVLFEECSQPALMKIKHVSNFVRKYSDTIAELQELQPSAKDFEVRSLVGCGHFAEVQVVREKATGDIYAMKVMKKKALLAQEQVSFFEEERNILSRSTSPWIPQLQYAFQDKNHLYLVMEYQPGGDLLSLLNRYEDQLDENLIQFYLAELILAVHSVHLMGYVHRDIKPENILVDRTGHIKLVDFGSAAKMNSNKMVNAKLPIGTPDYMAPEVLTVMNGDGKGTYGLDCDWWSVGVIAYEMIYGRSPFAEGTSARTFNNIMNFQRFLKFPDDPKVSSDFLDLIQSLLCGQKERLKFEGLCCHPFFSKIDWNNIRNSPPPFVPTLKSDDDTSNFDEPEKNSWVSSSPCQLSPSGFSGEELPFVGFSYSKALGILGRSESVVSGLDSPAKTSSMEKKLLIKSKELQDSQDKCHKMEQEMTRLHRRVSEVEAVLSQKEVELKASETQRSLLEQDLATYITECSSLKRSLEQARMEVSQEDDKALQLLHDIREQSRKLQEIKEQEYQAQVEEMRLMMNQLEEDLVSARRRSDLYESELRESRLAAEEFKRKATECQHKLLKAKDQGKPEVGEYAKLEKINAEQQLKIQELQEKLEKAVKASTEATELLQNIRQAKERAERELEKLQNREDSSEGIRKKLVEAEELEEKHREAQVSAQHLEVHLKQKEQHYEEKIKVLDNQIKKDLADKETLENMMQRHEEEAHEKGKILSEQKAMINAMDSKIRSLEQRIVELSEANKLAANSSLFTQRNMKAQEEMISELRQQKFYLETQAGKLEAQNRKLEEQLEKISHQDHSDKNRLLELETRLREVSLEHEEQKLELKRQLTELQLSLQERESQLTALQAARAALESQLRQAKTELEETTAEAEEEIQALTAHRDEIQRKFDALRNSCTVITDLEEQLNQLTEDNAELNNQNFYLSKQLDEASGANDEIVQLRSEVDHLRREITEREMQLTSQKQTMEALKTTCTMLEEQVMDLEALNDELLEKERQWEAWRSVLGDEKSQFECRVRELQRMLDTEKQSRARADQRITESRQVVELAVKEHKAEILALQQALKEQKLKAESLSDKLNDLEKKHAMLEMNARSLQQKLETERELKQRLLEEQAKLQQQMDLQKNHIFRLTQGLQEALDRADLLKTERSDLEYQLENIQVLYSHEKVKMEGTISQQTKLIDFLQAKMDQPAKKKKGLFSRRKEDPALPTQVPLQYNELKLALEKEKARCAELEEALQKTRIELRSAREEAAHRKATDHPHPSTPATARQQIAMSAIVRSPEHQPSAMSLLAPPSSRRKESSTPEEFSRRLKERMHHNIPHRFNVGLNMRATKCAVCLDTVHFGRQASKCLECQVMCHPKCSTCLPATCGLPAEYATHFTEAFCRDKMNSPGLQTKEPSSSLHLEGWMKVPRNNKRGQQGWDRKYIVLEGSKVLIYDNEAREAGQRPVEEFELCLPDGDVSIHGAVGASELANTAKADVPYILKMESHPHTTCWPGRTLYLLAPSFPDKQRWVTALESVVAGGRVSREKAEADAKLLGNSLLKLEGDDRLDMNCTLPFSDQVVLVGTEEGLYALNVLKNSLTHVPGIGAVFQIYIIKDLEKLLMIAGEERALCLVDVKKVKQSLAQSHLPAQPDISPNIFEAVKGCHLFGAGKIENGLCICAAMPSKVVILRYNENLSKYCIRKEIETSEPCSCIHFTNYSILIGTNKFYEIDMKQYTLEEFLDKNDHSLAPAVFAASSNSFPVSIVQVNSAGQREEYLLCFHEFGVFVDSYGRRSRTDDLKWSRLPLAFAYREPYLFVTHFNSLEVIEIQARSSAGTPARAYLDIPNPRYLGPAISSGAIYLASSYQDKLRVICCKGNLVKESGTEHHRGPSTSRSSPNKRGPPTYNEHITKRVASSPAPPEGPSHPREPSTPHRYREGRTELRRDKSPGRPLEREKSPGRMLSTRRERSPGRLFEDSSRGRLPAGAVRTPLSQVNKVWDQSSV'

# 对药物和靶标进行编码
d_out, mask_d_out = drug_encoder(d1)
# t_out, mask_t_out = target_encoder(t1)

# 转换为Tensor格式
d_out = paddle.to_tensor(d_out).unsqueeze(0)
# # t_out = paddle.to_tensor(t_out).unsqueeze(0)
mask_d_out = paddle.to_tensor(mask_d_out).unsqueeze(0)
# # mask_t_out = paddle.to_tensor(mask_t_out).unsqueeze(0)


# 构建输入数据
# input_data = [[d_out,t_out,mask_d_out,mask_t_out]]

# 定义模型
model_config = json.load(open('./config.json', 'r'))
model = MolTransModel(model_config)
model = model.cuda()


# 加载训练好的模型参数
model_state_dict = paddle.load('DAVIS_bestAUC_model_cls1')
model.set_state_dict(model_state_dict)
# model.set_state_dict(paddle.load('DAVIS_bestAUC_model_cls1-20'))

# 设置损失函数
sig = paddle.nn.Sigmoid()

# 预测
model.eval()

with paddle.no_grad():
    for i, sequence in enumerate(protein_sequences):
        # 对蛋白质序列进行编码
        t_out, mask_t_out = target_encoder(sequence)
        t_out = paddle.to_tensor(t_out).unsqueeze(0)
        mask_t_out = paddle.to_tensor(mask_t_out).unsqueeze(0)
        
        
        output=model(d_out,t_out,mask_d_out,mask_t_out)
        interaction=paddle.squeeze(sig(output))
        interaction=interaction.detach().cpu().numpy()
        print(f"Prediction for protein {i+1}: {interaction}")
# print(interaction)